"""
LLM4CP工具函数文件
包含RMSE计算、损失函数等工具函数
"""

import torch
import torch.nn.functional as F


def calculate_matlab_style_rmse(pred, target):
    """改进的MATLAB风格RMSE计算"""
    # 确保维度匹配
    if pred.shape != target.shape:
        min_dim = min(pred.shape[-1], target.shape[-1])
        pred = pred[..., :min_dim]
        target = target[..., :min_dim]
    
    diff = pred - target
    squared_diff = torch.abs(diff) ** 2
    sum_squared_diff = torch.sum(squared_diff, dim=2)
    N1 = pred.shape[-1]  # 动态特征数量
    mse_per_time = sum_squared_diff / N1
    rmse_per_time = torch.sqrt(mse_per_time + 1e-12)  # 更小的epsilon
    
    # 加权平均 - 给后面的时间步更高权重
    weights = torch.linspace(1.0, 1.5, steps=rmse_per_time.shape[1]).to(pred.device)
    weighted_rmse = torch.sum(rmse_per_time * weights, dim=1) / torch.sum(weights)
    rmse = torch.mean(weighted_rmse)
    
    return rmse


def calculate_consistency_loss(spatial_pred, channel_pred):
    """计算空间-信道一致性损失"""
    # 提取空间预测的前96维与信道预测比较
    spatial_channel_part = spatial_pred[..., :96]
    consistency_loss = F.mse_loss(spatial_channel_part, channel_pred)
    return consistency_loss


def calculate_smoothness_loss(pred):
    """计算平滑性损失 - 鼓励预测的平滑性"""
    # 时间维度的平滑性
    time_diff = torch.diff(pred, dim=1)
    time_smoothness = torch.mean(torch.pow(time_diff, 2))
    
    # 特征维度的平滑性
    feature_diff = torch.diff(pred, dim=2)
    feature_smoothness = torch.mean(torch.pow(feature_diff, 2))
    
    return time_smoothness + feature_smoothness * 0.5 